clc;
clear;

add_path;

% H = 288;
% W = 352;
H = 144;
W = 176;
st = 0;
dt = 1;
et = 249;
tau = floor((et - st) / dt) + 1;
C = 3;

CPALS_MaxIter = 50;

o_images = zeros(H, W, tau, C);

video_path = '.\data\';
video_name = 'lift'; % akiyo  fountain  football  highway  mobile  lift  smoke  self_01  self_02  self_03  self_04  self_05
video_size ='176x144'; % 176x144  352x288

% Ns = [30, 15, 10, 6, 5, 3, 2, 1];
% Rs = [1, 2, 3, 5, 6, 10, 15, 30];
% Ns = [18, 9, 6, 3, 2, 1];
% Rs = [1, 2, 3, 6, 9, 18];
Ns = [2, 1];
Rs = [9, 18];
K = length(Ns);

multi_iter = Ns(1);
R = Rs(1);

for k = 1 : K
    
    multi_iter = Ns(k);
    R = Rs(k);
    psnrs = zeros(multi_iter, tau + 2);
    
    i = 0;
    for t = st : dt : et
        image = imread([video_path video_name '\' video_size '\' num2str(t, '%03d') '.bmp']);
        i = i + 1;
        o_images(:, :, i, :) = image;
    end
    
    X = tensor(o_images);
    Z = X;
    
    lambda = zeros(multi_iter * R, 1);
    U1 = zeros(H, multi_iter * R);
    U2 = zeros(W, multi_iter * R);
    U3 = zeros(tau, multi_iter * R);
    U4 = zeros(C, multi_iter * R);
    
    dec_coef_matfile = ['.\result\' video_name '_comp_MultiRankR_coefficients_(st,dt,et,tau)=(' num2str(st, '%03d') ',' num2str(dt, '%03d') ',' num2str(et, '%03d') ',' num2str(tau, '%03d') ')_MultiIter=' num2str(multi_iter, '%03d') '_R=' num2str(R, '%03d') '_CPALSMaxIter=' num2str(CPALS_MaxIter, '%03d') '.mat'];
    if exist(dec_coef_matfile, 'file')
        load(dec_coef_matfile);
    else
        for n = 1 : multi_iter
            
            fprintf('\ndec, k=%d/%d, n=%d/%d\n', k, K, n, multi_iter);
            
            if n > 1
                Z = Z .* 0;
                for r = 1:R
                    ktr = ktensor(lambda(index(r)), {U1(:, index(r)), U2(:, index(r)), U3(:, index(r)), U4(:, index(r))});
                    Z = Z + full(ktr);
                end
                X = X - Z;
            end
            
            Y = cp_als(X, R, 'tol', 1e-12, 'maxiters', CPALS_MaxIter);
            
            index = (n - 1) * R + 1 : n * R;
            lambda(index) = Y.lambda;
            U1(:, index) = Y.U{1};
            U2(:, index) = Y.U{2};
            U3(:, index) = Y.U{3};
            U4(:, index) = Y.U{4};
            
        end
        save(dec_coef_matfile, 'lambda', 'U1', 'U2', 'U3', 'U4');
    end
    
    dec_coef_datfile = ['.\result\' video_name '_comp_MultiRankR_coefficients_(st,dt,et,tau)=(' num2str(st, '%03d') ',' num2str(dt, '%03d') ',' num2str(et, '%03d') ',' num2str(tau, '%03d') ')_MultiIter=' num2str(multi_iter, '%03d') '_R=' num2str(R, '%03d') '_CPALSMaxIter=' num2str(CPALS_MaxIter, '%03d') '.dat'];
    if ~exist(dec_coef_datfile, 'file')
        f = fopen(dec_coef_datfile, 'wb');
        fwrite(f, lambda, 'double');
        fwrite(f, U1, 'double');
        fwrite(f, U2, 'double');
        fwrite(f, U3, 'double');
        fwrite(f, U4, 'double');
        fclose(f);
    end
    
    dec_coef_q_matfile = ['.\result\' video_name '_comp_MultiRankR_coefficients_q_(st,dt,et,tau)=(' num2str(st, '%03d') ',' num2str(dt, '%03d') ',' num2str(et, '%03d') ',' num2str(tau, '%03d') ')_MultiIter=' num2str(multi_iter, '%03d') '_R=' num2str(R, '%03d') '_CPALSMaxIter=' num2str(CPALS_MaxIter, '%03d') '.mat'];
    if exist(dec_coef_q_matfile, 'file')
        load(dec_coef_q_matfile);
    else
        allU_char = zeros(W + H + C + tau + 1, R * multi_iter);
        allU_char(1 : W, :) = U2;
        allU_char((W + 1) : (W + H), :) = U1;
        allU_char((W + H + 1) : (W + H + C), :) = U4;
        allU_char((W + H + C + 1) : (W + H + C + tau), :) = U3;
        mm_allU = [min(allU_char(:)) max(allU_char(:))];
        allU_char = uint8(mat2gray(allU_char) * 255);
        
        mm_lambda = [min(lambda) max(lambda)];
        lambda_char = uint8(mat2gray(lambda) * 255);
        
        save(dec_coef_q_matfile, 'mm_allU', 'allU_char', 'mm_lambda', 'lambda_char');
    end
    allU_double = double(allU_char) / 255 * (mm_allU(2) - mm_allU(1)) + mm_allU(1);
    lambda_double = double(lambda_char) / 255 * (mm_lambda(2) - mm_lambda(1)) + mm_lambda(1);
    
    dec_coef_q_datfile = ['.\result\' video_name '_comp_MultiRankR_coefficients_q_(st,dt,et,tau)=(' num2str(st, '%03d') ',' num2str(dt, '%03d') ',' num2str(et, '%03d') ',' num2str(tau, '%03d') ')_MultiIter=' num2str(multi_iter, '%03d') '_R=' num2str(R, '%03d') '_CPALSMaxIter=' num2str(CPALS_MaxIter, '%03d') '.dat'];
    if ~exist(dec_coef_q_datfile, 'file')
        f = fopen(dec_coef_q_datfile, 'wb');
        fwrite(f, mm_allU, 'double');
        fwrite(f, allU_char, 'uint8');
        fwrite(f, mm_lambda, 'double');
        fwrite(f, lambda_char, 'uint8');
        fclose(f);
    end
    
    U2 = allU_double(1 : W, :);
    U1 = allU_double((W + 1) : (W + H), :);
    U4 = allU_double((W + H + 1) : (W + H + C), :);
    U3 = allU_double((W + H + C + 1) : (W + H + C + tau), :);
    lambda = lambda_double;
    
    Z = Z .* 0;
    for n = 1 : multi_iter
        
        fprintf('\nrec, k=%d/%d, n=%d/%d\n', k, K, n, multi_iter);
        
        index = (n - 1) * R + 1 : n * R;
        for r = 1:R
            ktr = ktensor(lambda(index(r)), {U1(:, index(r)), U2(:, index(r)), U3(:, index(r)), U4(:, index(r))});
            Z = Z + full(ktr);
        end
        
        i = 0;
        for t = st : dt : et
            i = i + 1;
            image(:, :, 1) = uint8(Z.data(:, :, i, 1) + 0.5);
            image(:, :, 2) = uint8(Z.data(:, :, i, 2) + 0.5);
            image(:, :, 3) = uint8(Z.data(:, :, i, 3) + 0.5);
            psnrs(n, i) = psnr(o_images(:, :, i, :), image);
%             imshow(image);
%             title(['R=' num2str(R) ', n=' num2str(n) '/' num2str(multi_iter) ', t=' num2str(t) '/' num2str(et) ', PSNR=' num2str(psnrs(n, i))]);
%             pause(1/25);
            imwrite(image, [video_path video_name '\comp\MultiRankR_(st,dt,et,tau)=(' num2str(st, '%03d') ',' num2str(dt, '%03d') ',' num2str(et, '%03d') ',' num2str(tau, '%03d') ')_MultiIter=' num2str(multi_iter, '%03d') '_R=' num2str(R, '%03d') '_CPALSMaxIter=' num2str(CPALS_MaxIter, '%03d') '_n=' num2str(n, '%03d') '_t=' num2str(t, '%03d') '.png']);
        end
        psnrs(n, tau + 1) = mean(psnrs(n, 1 : tau));
        psnrs(n, tau + 2) = n;
        
    end
    
    save(['.\result\' video_name '_comp_MultiRankR_psnrs_(st,dt,et,tau)=(' num2str(st, '%03d') ',' num2str(dt, '%03d') ',' num2str(et, '%03d') ',' num2str(tau, '%03d') ')_MultiIter=' num2str(multi_iter, '%03d') '_R=' num2str(R, '%03d') '_CPALSMaxIter=' num2str(CPALS_MaxIter, '%03d') '.mat'], 'psnrs');
    
end
